import cvxpy as cp
import numpy as np

from src.drl.a2c import A2C


class RA:

    @classmethod
    def a2c(cls, env):
        state_dim = env.config.dnn_dim
        return A2C(state_dim=state_dim, action_dim=env.M, action_prob_dim=env.M,
                   lr_a=0.0001, lr_c=0.0001,
                   epsilon=env.config.dnn_epsilon, batch_size=env.config.dnn_batch_size)

    @classmethod
    def sca(cls, env):
        sca_model = SCA(env)
        return sca_model


class SCA:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = self.init()
        self.action_bi = self.action[:].astype(int)
        self.bandwidth_bi, self.frequency_bi = self.bandwidth, self.frequency

    def run(self):
        # 1. 松弛为连续变量并计算最优值
        for i in range(self.env.config.sca_episode):
            self.optimize_v2()
            if i % 10 == 0:
                print(f"\tSCA: {i}, avg_action: {np.average(self.action)}, avg_bandwidth: {np.average(self.bandwidth)}, avg_frequency: {np.average(self.frequency)}")
        # 2. 将连续变量离散化（实验：并计算离散卸载决策下的资源分配最优方案，验证该操作对实验结果的影响）
        self.action_bi = self.action.squeeze()
        self.bandwidth_bi = self.bandwidth.squeeze()
        self.frequency_bi = self.frequency.squeeze()
        # 3. 返回离散卸载决策对应的最优资源分配结果
        return self.action_bi, self.bandwidth_bi, self.frequency_bi

    def init(self):
        return (np.ones(self.env.M).reshape(-1, 1),
                np.ones(self.env.M).reshape(-1, 1),
                np.ones(self.env.M).reshape(-1, 1))

    def discrete_half(self):
        self.action_bi = [1 if a >= 0.5 else 0 for a in self.action]

    def discrete_sequence(self):
        action_bi = np.zeros(self.env.M).astype(int)
        offloading_num = int(np.mean(self.action) * 100)
        idx = np.argsort(self.action.squeeze())[0:offloading_num]
        action_bi[idx] = 1
        self.action_bi = action_bi.tolist()

    def optimize(self):
        # 1. 初始化变量和相关参数
        drone = self.env.drones[0]
        drone_pos = drone.pos
        action_opt = np.zeros(self.env.M).reshape(-1, 1)
        bandwidth_opt = np.zeros(self.env.M).reshape(-1, 1)
        frequency_opt = np.zeros(self.env.M).reshape(-1, 1)
        # 2. 顺序解决 G 个凸优化问题
        time_wait_drone = 0
        for part_idx in self.env.path:
            # 2.1. 获取 PT 的信息，计算等待时间
            part = self.env.parts[part_idx]
            ues_num = len(part.ues)
            distance = np.linalg.norm(drone_pos - part.pos).astype(int)
            time_wait_drone += np.round(distance / drone.speed, decimals=2)
            drone_pos = drone.update_pos(part.pos)
            # 2.2. 创建优化变量
            action_part = cp.Variable(ues_num, nonneg=True)
            bandwidth_part = cp.Variable(ues_num, nonneg=True)
            frequency_part = cp.Variable(ues_num, nonneg=True)
            # 2.3. 计算常量参数
            ues_matrix = np.array([[ue.idx, ue.trans_power, ue.task.size, ue.task.cycles_per_bit,
                                    ue.task.weight_time, ue.task.weight_energy, ue.t_local, ue.e_local] for ue in part.ues])
            ues_pos = np.array([[ue.pos[0], ue.pos[1], ue.pos[2]] for ue in part.ues])
            ues_idx = ues_matrix[:, 0].astype(int)
            channel_gain = self.env.config.channel_gain / (np.linalg.norm(drone_pos - ues_pos, axis=1).astype(int) ** 2)
            constants = np.hstack((channel_gain[:, np.newaxis], ues_matrix[:, 1:]))
            variables = np.concatenate((self.action[ues_idx, :], self.bandwidth[ues_idx, :], self.frequency[ues_idx, :]), axis=1)
            value, derivative, hessian = self.sca_f(constants, variables)
            d1 = cp.multiply(action_part - variables[:, 0], derivative[:, 0])
            d2 = cp.multiply(bandwidth_part - variables[:, 1], derivative[:, 1])
            d3 = cp.multiply(frequency_part - variables[:, 2], derivative[:, 2])
            # 2.4. 创建约束和优化目标（一阶 Taylor 展开）
            obj = cp.Minimize(value + cp.sum(cp.sum([d1, d2, d3])))
            constraints = [action_part <= self.env.config.sca_vars_bound[0, 1],
                           bandwidth_part <= self.env.config.sca_vars_bound[1, 1],
                           frequency_part <= self.env.config.sca_vars_bound[2, 1],
                           cp.sum(bandwidth_part) <= self.env.config.sca_vars_bound[1, 1],
                           cp.sum(frequency_part) <= self.env.config.sca_vars_bound[2, 1]]
            for i in range(ues_num):
                constraints.append(action_part[i] >= self.env.config.sca_vars_bound[0, 0])
                constraints.append(bandwidth_part[i] >= self.env.config.sca_vars_bound[1, 0])
                constraints.append(frequency_part[i] >= self.env.config.sca_vars_bound[2, 0])
            prob = cp.Problem(objective=obj, constraints=constraints)
            prob.solve()
            # prob.solve(solver='OSQP', max_iter=10000)
            # 2.5. 保存当前 Part 的最优资源分配策略
            for i, idx in enumerate(ues_idx):
                action_opt[int(idx)] = action_part.value[i]
                bandwidth_opt[int(idx)] = bandwidth_part.value[i]
                frequency_opt[int(idx)] = frequency_part.value[i]
            # 2.6. 更新等待时间
            trans_rate = cp.multiply(bandwidth_part,
                                     np.log2(1 + ues_matrix[:, 1] * channel_gain / self.env.config.channel_noise))
            time_off = cp.multiply(ues_matrix[:, 2], cp.inv_pos(trans_rate))
            time_exe = cp.multiply(ues_matrix[:, 2] * ues_matrix[:, 3], cp.inv_pos(frequency_part))
            time_wait_drone += cp.max(cp.sum([time_off, time_exe])).value
        self.action += self.env.config.sca_vars_step[0] * (action_opt - self.action)
        self.bandwidth += self.env.config.sca_vars_step[1] * (bandwidth_opt - self.bandwidth)
        self.frequency += self.env.config.sca_vars_step[2] * (frequency_opt - self.frequency)
        self.bandwidth_bi = self.bandwidth.squeeze()
        self.frequency_bi = self.frequency.squeeze()

    def optimize_v2(self):
        # 1. 初始化变量和相关参数
        drone = self.env.drones[0]
        action_opt = np.zeros(self.env.M)
        bandwidth_opt = np.zeros(self.env.M)
        frequency_opt = np.zeros(self.env.M)

        fm1 = []
        fm2 = []
        fm3 = []
        fm4 = []
        action_var = cp.Variable(self.env.M, nonneg=True)
        bandwidth_var = cp.Variable(self.env.M, nonneg=True)
        frequency_var = cp.Variable(self.env.M, nonneg=True)
        constraints = []
        for ue in self.env.ues:
            k1 = ue.task.size / np.log2(1 + ue.trans_power * self.env.config.channel_gain / self.env.config.channel_noise)
            k2 = ue.task.size * ue.task.cycles_per_bit
            k3 = k2 / ue.cpu_frequency
            f1_0 = self.action[ue.idx] * k1 / self.bandwidth[ue.idx]
            f1_partial_a = k1 / self.bandwidth[ue.idx]
            f1_partial_a = cp.multiply(f1_partial_a, action_var[ue.idx] - self.action[ue.idx])
            f1_partial_b = (-self.action[ue.idx] * k1) / (self.bandwidth[ue.idx] ** 2)
            f1_partial_b = cp.multiply(f1_partial_b, bandwidth_var[ue.idx] - self.bandwidth[ue.idx])
            fm1.append(f1_0 + f1_partial_a + f1_partial_b)
            f2_0 = self.action[ue.idx] * k2 / self.frequency[ue.idx]
            f2_partial_a = k2 / self.frequency[ue.idx]
            f2_partial_a = cp.multiply(f2_partial_a, action_var[ue.idx] - self.action[ue.idx])
            f2_partial_f = (-self.action[ue.idx] * k2) / (self.frequency[ue.idx] ** 2)
            f2_partial_f = cp.multiply(f2_partial_f, frequency_var[ue.idx] - self.frequency[ue.idx])
            fm2.append(f2_0 + f2_partial_a + f2_partial_f)
            f3 = cp.multiply(k3, 1 - action_var[ue.idx])
            fm3.append(f3)
            f4_0 = (1 - self.action[ue.idx]) * k1 / self.bandwidth[ue.idx]
            f4_partial_a = -k1 / self.bandwidth[ue.idx]
            f4_partial_a = cp.multiply(f4_partial_a, action_var[ue.idx] - self.action[ue.idx])
            f4_partial_b = (self.action[ue.idx] - 1) * k1 / (self.bandwidth[ue.idx] ** 2)
            f4_partial_b = cp.multiply(f4_partial_b, bandwidth_var[ue.idx] - self.bandwidth[ue.idx])
            fm4.append(f4_0 + f4_partial_a + f4_partial_b)
            # Constraint
            constraints.append(bandwidth_var[ue.idx] >= self.env.config.sca_vars_bound[1, 0])
            constraints.append(frequency_var[ue.idx] >= self.env.config.sca_vars_bound[2, 0])
            constraints.append(action_var[ue.idx] <= self.env.config.sca_vars_bound[0, 1])
            constraints.append(bandwidth_var[ue.idx] <= self.env.config.sca_vars_bound[1, 1])
            constraints.append(frequency_var[ue.idx] <= self.env.config.sca_vars_bound[2, 1])

        t_part_wait = []
        fm1_np = np.array(fm1)
        fm2_np = np.array(fm2)
        for i in range(self.env.G):
            drone_pos = self.env.parts[self.env.path[0]].pos
            t_pre_part_service = []
            for j in range(i):
                part_idx = self.env.path[j]
                part = self.env.parts[part_idx]
                part_ues_idx = self.env.parts_ues[part_idx]
                t_fly = np.linalg.norm(drone_pos - part.pos).astype(int) / drone.speed
                fm1_part = fm1_np[part_ues_idx].tolist()
                fm2_part = fm2_np[part_ues_idx].tolist()
                t_service = [a + b for a, b in zip(fm1_part, fm2_part)]
                t_pre_part_service.append(t_fly + cp.max(cp.vstack(t_service)))
            t_part_wait.append(cp.sum(t_pre_part_service))
            # Cons
            ues_idx = self.env.parts_ues[self.env.path[i]]
            bandwidth_part = bandwidth_var[ues_idx]
            frequency_part = frequency_var[ues_idx]
            constraints.append(cp.sum(bandwidth_part) <= self.env.config.sca_vars_bound[1, 1])
            constraints.append(cp.sum(frequency_part) <= self.env.config.sca_vars_bound[2, 1])

        ues_cost = []
        for ue in self.env.ues:
            # Cost
            t_wait = t_part_wait[ue.part_seq]
            t_fly = 0 if ue.part_seq == 0 else (np.linalg.norm(
                self.env.parts[self.env.path[ue.part_seq-1]].pos - self.env.parts[ue.part].pos).astype(int) / drone.speed)
            t_service = fm1[ue.idx] + fm2[ue.idx]
            ues_cost.append(cp.maximum(t_wait + t_fly + t_service, fm3[ue.idx] + fm4[ue.idx]))

        obj = cp.Minimize(cp.sum(ues_cost))
        prob = cp.Problem(objective=obj, constraints=constraints)
        prob.solve()
        # 4. 获取最优结果
        for ue in self.env.ues:
            action_opt[ue.idx] = action_var[ue.idx].value
            bandwidth_opt[ue.idx] = bandwidth_var[ue.idx].value
            frequency_opt[ue.idx] = frequency_var[ue.idx].value
        # 5. 处理结果格式
        self.action += self.env.config.sca_vars_step[0] * (action_opt - self.action)
        self.bandwidth += self.env.config.sca_vars_step[1] * (bandwidth_opt - self.bandwidth)
        self.frequency += self.env.config.sca_vars_step[2] * (frequency_opt - self.frequency)

    def optimize_v3(self):
        self.action = self.action.squeeze()
        self.bandwidth = self.bandwidth.squeeze()
        self.frequency = self.frequency.squeeze()

        # 1. 初始化变量和相关参数
        drone = self.env.drones[0]
        action_opt = np.zeros(self.env.M)
        bandwidth_opt = np.zeros(self.env.M)
        frequency_opt = np.zeros(self.env.M)
        action_var = cp.Variable(self.env.M, nonneg=True)
        bandwidth_var = cp.Variable(self.env.M, nonneg=True)
        frequency_var = cp.Variable(self.env.M, nonneg=True)
        constraints = []
        constraints.append(bandwidth_var >= self.env.config.sca_vars_bound[1, 0])
        constraints.append(frequency_var >= self.env.config.sca_vars_bound[2, 0])
        constraints.append(action_var <= self.env.config.sca_vars_bound[0, 1])
        constraints.append(bandwidth_var <= self.env.config.sca_vars_bound[1, 1])
        constraints.append(frequency_var <= self.env.config.sca_vars_bound[2, 1])

        ues_matrix = np.array([[ue.idx, ue.trans_power, ue.cpu_frequency, ue.task.size, ue.task.cycles_per_bit, ue.task.weight_time, ue.task.weight_energy] for ue in self.env.ues])
        ue_trans_power = ues_matrix[:, 1].squeeze()
        ue_cpu_frequency = ues_matrix[:, 2].squeeze()
        ue_task_size = ues_matrix[:, 3].squeeze()
        ue_task_cpu = ues_matrix[:, 4].squeeze()

        k1 = ue_task_size / np.log2(1 + ue_trans_power * self.env.config.channel_gain / self.env.config.channel_noise)
        k2 = ue_task_size * ue_task_cpu
        k3 = k2 / ue_cpu_frequency

        f1_0 = self.action * k1 / self.bandwidth
        f1_partial_a = k1 / self.bandwidth
        f1_partial_b = (-self.action * k1) / (self.bandwidth ** 2)
        fm1 = f1_0 + cp.multiply(f1_partial_a, action_var - self.action) + cp.multiply(f1_partial_b, bandwidth_var - self.bandwidth)

        f2_0 = self.action * k2 / self.frequency
        f2_partial_a = k2 / self.frequency
        f2_partial_f = (-self.action * k2) / (self.frequency ** 2)
        fm2 = f2_0 + cp.multiply(f2_partial_a, action_var - self.action) + cp.multiply(f2_partial_f, frequency_var - self.frequency)

        fm3 = cp.multiply(k3, 1 - action_var)

        f4_0 = (1 - self.action) * k1 / self.bandwidth
        f4_partial_a = -k1 / self.bandwidth
        f4_partial_b = (self.action - 1) * k1 / (self.bandwidth ** 2)
        fm4 = f4_0 + cp.multiply(f4_partial_a, action_var - self.action) + cp.multiply(f4_partial_b, bandwidth_var - self.bandwidth)

        t_part_wait = []
        for i in range(self.env.G):
            drone_pos = self.env.parts[self.env.path[0]].pos
            t_pre_part_service = []
            for j in range(i):
                part_idx = self.env.path[j]
                part = self.env.parts[part_idx]
                part_ues_idx = self.env.parts_ues[part_idx]
                t_fly = np.linalg.norm(drone_pos - part.pos).astype(int) / drone.speed
                fm1_part = fm1[part_ues_idx]
                fm2_part = fm2[part_ues_idx]
                t_service = [a + b for a, b in zip(fm1_part, fm2_part)]
                t_pre_part_service.append(t_fly + cp.max(cp.vstack(t_service)))
            t_part_wait.append(cp.sum(t_pre_part_service))
            # Cons
            ues_idx = self.env.parts_ues[self.env.path[i]]
            bandwidth_part = bandwidth_var[ues_idx]
            frequency_part = frequency_var[ues_idx]
            constraints.append(cp.sum(bandwidth_part) <= self.env.config.sca_vars_bound[1, 1])
            constraints.append(cp.sum(frequency_part) <= self.env.config.sca_vars_bound[2, 1])

        t_wait = []
        for ue in self.env.ues:
            t_fly = 0 if ue.part_seq == 0 else (np.linalg.norm(self.env.parts[self.env.path[ue.part_seq - 1]].pos - self.env.parts[ue.part].pos).astype(int) / drone.speed)
            t_wait.append(t_part_wait[ue.part_seq] + t_fly)
        t_wait = cp.vstack(t_wait)
        t_wait_r = cp.reshape(t_wait, (t_wait.shape[0],))
        le = t_wait_r + fm1 + fm2
        re = fm3 + fm4
        ues_cost = cp.maximum(le, re)

        obj = cp.Minimize(cp.sum(ues_cost))
        prob = cp.Problem(objective=obj, constraints=constraints)
        prob.solve()
        # 4. 获取最优结果
        for ue in self.env.ues:
            action_opt[ue.idx] = action_var[ue.idx].value
            bandwidth_opt[ue.idx] = bandwidth_var[ue.idx].value
            frequency_opt[ue.idx] = frequency_var[ue.idx].value
        # 5. 处理结果格式
        self.action += self.env.config.sca_vars_step[0] * (action_opt - self.action)
        self.bandwidth += self.env.config.sca_vars_step[1] * (bandwidth_opt - self.bandwidth)
        self.frequency += self.env.config.sca_vars_step[2] * (frequency_opt - self.frequency)

    def discrete_optimize(self):
        self.env.execute(action=self.action_bi)
        # 1. 初始化变量和相关参数
        drone = self.env.drones[0]
        drone_pos = drone.pos
        bandwidth = np.zeros(self.env.M)
        frequency = np.zeros(self.env.M)
        # 2. 顺序解决 G 个凸优化问题
        time_wait_drone = 0
        for part_idx in self.env.path:
            # 2.0. 前置检查
            part = self.env.parts[part_idx]
            ue_idx_loc = np.array([ue.idx for ue in part.ues if not ue.offloading], dtype=int)
            ue_idx_off = np.array([ue.idx for ue in part.ues if ue.offloading], dtype=int)
            # 2.0.1. 只有一个 UD 需要资源分配
            if len(ue_idx_off) == 1:
                bandwidth[ue_idx_off[0]] = self.env.config.uav_bandwidth
                frequency[ue_idx_off[0]] = self.env.config.uav_cpu_frequency
                continue
            # 2.0.2. 需要本地计算的 UDs 不参与资源分配
            if len(ue_idx_loc):
                bandwidth[ue_idx_loc] = 0
                frequency[ue_idx_loc] = 0
            # 2.0.3. 检查参与资源分配的 UDs 数目
            ues_num = len(ue_idx_off)
            if ues_num == 0:
                continue
            # 2.1. 获取 PT 的信息，计算等待时间
            distance = np.linalg.norm(drone_pos - part.pos).astype(int)
            time_wait_drone += np.round(distance / drone.speed, decimals=2)
            drone_pos = part.pos
            # 2.2. 创建优化变量
            bandwidth_part = cp.Variable(ues_num, nonneg=True)
            frequency_part = cp.Variable(ues_num, nonneg=True)
            # 2.3. 定义优化问题
            ues_matrix = np.array([[ue.idx, ue.trans_power, ue.task.size, ue.task.cycles_per_bit,
                                    ue.task.weight_time, ue.task.weight_energy] for ue in part.ues if ue.offloading])
            ues_pos = np.array([ue.pos for ue in part.ues if ue.offloading])
            channel_gain = self.env.config.channel_gain / (np.linalg.norm(drone_pos - ues_pos, axis=1).astype(int) ** 2)
            trans_power = ues_matrix[:, 1]
            trans_rate = cp.multiply(bandwidth_part, np.log2(1 + trans_power * channel_gain / self.env.config.channel_noise))
            time_off = cp.multiply(ues_matrix[:, 2], cp.inv_pos(trans_rate))
            time_exe = cp.multiply(ues_matrix[:, 2] * ues_matrix[:, 3], cp.inv_pos(frequency_part))
            lambda_t = ues_matrix[:, 4]
            lambda_e = ues_matrix[:, 5]
            time_total = cp.sum([time_wait_drone, time_off, time_exe])
            energy_total = cp.multiply(trans_power, time_off)
            # 2.4. 创建约束和优化目标
            obj = cp.Minimize(cp.sum(cp.multiply(lambda_t, time_total) + cp.multiply(lambda_e, energy_total)))
            constraints = [cp.sum(bandwidth_part) <= self.env.config.sca_vars_bound[1, 1],
                           cp.sum(frequency_part) <= self.env.config.sca_vars_bound[2, 1]]
            for i in range(ues_num):
                constraints.append(bandwidth_part[i] >= 1)
                constraints.append(frequency_part[i] >= 100)
            prob = cp.Problem(objective=obj, constraints=constraints)
            prob.solve()
            # 2.5. 保存当前 Part 的最优资源分配策略
            for i, idx in enumerate(ues_matrix[:, 0]):
                bandwidth[int(idx)] = bandwidth_part.value[i]
                frequency[int(idx)] = frequency_part.value[i]
            # 2.6. 更新等待时间
            time_wait_drone += cp.max(cp.sum([time_off, time_exe])).value
        self.bandwidth_bi = bandwidth
        self.frequency_bi = frequency

    def sca_f(self, constants, variables):
        # 1. 常量参数处理
        ues_num = constants.shape[0]
        ues_prop_num = constants.shape[1]
        channel_gain, trans_power, task_size, task_cpu, lambda_t, lambda_e, time_local, energy_local = np.hsplit(
            constants, ues_prop_num)
        action, bandwidth, frequency = np.hsplit(variables, variables.shape[1])
        k1 = task_size / np.log2(1 + trans_power * channel_gain / self.env.config.channel_noise)
        k2 = task_size * task_cpu
        k3 = time_local
        k4 = k1 * trans_power
        k5 = energy_local
        # 2. 计算函数在 variables 处的函数值 (Real number)
        value = (lambda_t * (action * (k1 / bandwidth + k2 / frequency) + (1 - action) * k3) +
                 lambda_e * (action * k4 / bandwidth + (1 - action) * k5))
        value = np.sum(value)
        # 3. 计算函数在 variables 处的一阶偏导 (ues_num * 3)
        partial_a = lambda_t * (
                    k1 / bandwidth + k2 / frequency) - lambda_t * k3 + lambda_e * k4 / bandwidth - lambda_e * k5
        partial_b = -action * (lambda_t * k1 + lambda_e * k4) / (bandwidth * bandwidth)
        partial_f = -action * (lambda_t * k2) / (frequency * frequency)
        derivative = np.concatenate((partial_a, partial_b, partial_f), axis=1)
        # 4. 计算函数在 variables 处的 Hessian 矩阵 (ues_num * 3 * 3, hessian[0] = Hessian of ue-0)
        partial_aa = np.zeros(ues_num, dtype=float).reshape(-1, 1)
        partial_ab = -(lambda_t * k1 + lambda_e * k4) / (bandwidth * bandwidth)
        partial_af = -(lambda_t * k2) / (frequency * frequency)
        partial_ba = partial_ab
        partial_bb = 2 * action * (lambda_t * k1 + lambda_e * k4) / np.power(bandwidth, 3)
        partial_bf = np.zeros(ues_num, dtype=float).reshape(-1, 1)
        partial_fa = partial_af
        partial_fb = np.zeros(ues_num, dtype=float).reshape(-1, 1)
        partial_ff = 2 * action * (lambda_t * k2) / np.power(frequency, 3)
        hessian = np.transpose(np.concatenate((np.array([partial_aa, partial_ab, partial_af]),
                                               np.array([partial_ba, partial_bb, partial_bf]),
                                               np.array([partial_fa, partial_fb, partial_ff])), axis=2), (1, 0, 2))
        return value, derivative, hessian

    def sca_g(self, constants, variables):
        # 1. 常量参数处理
        ues_num = constants.shape[0]
        ues_prop_num = constants.shape[1]
        channel_gain, trans_power, task_size, task_cpu, lambda_t, lambda_e, time_local, energy_local = np.hsplit(
            constants, ues_prop_num)
        action, bandwidth, frequency = np.hsplit(variables, variables.shape[1])
        k1 = task_size / np.log2(1 + trans_power * channel_gain / self.env.config.channel_noise)
        k2 = task_size * task_cpu
        k3 = time_local
        # 2. 计算函数在 variables 处的函数值 (Real number)
        value = action * (k1 / bandwidth + k2 / frequency) + (1 - action) * k3
        # 3. 计算函数在 variables 处的一阶偏导 (ues_num * 3)
        partial_a = k1 / bandwidth + k2 / frequency
        partial_b = -action * k1 / (bandwidth * bandwidth)
        partial_f = -action * k2 / (frequency * frequency)
        derivative = np.concatenate((partial_a, partial_b, partial_f), axis=1)
        return value.squeeze(), derivative

